
from collections.abc import Iterable

import importlib
import logging
import lxml.etree
import math
import numba
import numpy as np


logger = logging.getLogger(__name__)


def dynamic_import(module_name, py_path):
    module_spec = importlib.util.spec_from_file_location(module_name, py_path)
    module = importlib.util.module_from_spec(module_spec)
    module_spec.loader.exec_module(module)
    return module


def flatten(arr):
    for e in arr:
        if not isinstance(e, Iterable) or isinstance(e, str):
            yield e
        else:
            yield from flatten(e)


@numba.jit(nopython=True)
def normalize(arr: np.ndarray):
    """to [0.0, 1.0]"""
    arr = arr.astype(np.float64)
    imin = np.nanmin(arr)
    imax = np.nanmax(arr)
    return (arr - imin) / (imax - imin)


@numba.jit(nopython=True)
def normalize_p(arr: np.ndarray, lp: float, hp: float):
    arr = arr.astype(np.float64)
    imin = np.percentile(arr, lp)
    imax = np.percentile(arr, hp)
    return np.clip((arr - imin) / (imax - imin), 0.0, 1.0)


@numba.jit(nopython=True)
def remap_range(arr: np.ndarray, low, high):
    return low + (high - low) * arr


# @numba.jit(nopython=True)
def remap_to_rgb(arr: np.ndarray, lp: float, hp: float):
    # imin = np.percentile(arr, lp)
    # imax = np.percentile(arr, hp)
    # print(imin, imax)

    arr = normalize_p(arr, lp, hp)

    # imin = np.nanmin(arr)
    # imax = np.nanmax(arr)
    # print(imin, imax)

    return remap_range(arr, 0.0, 255.0).astype(np.uint8)


def xyxy2axywh(xyxy):
    """四方 -> 左上+宽高"""
    xmin, ymin, xmax, ymax = xyxy
    w, h = xmax - xmin, ymax - ymin
    return (xmin, ymin, w, h)


def axywh2xyxy(axywh):
    """左上+宽高 -> 四方"""
    x, y, w, h = axywh
    return (x, y, x + w, y + h)


def xyxy2cxywh(xyxy):
    """四方 -> 中心+宽高"""
    xmin, ymin, xmax, ymax = xyxy
    w, h = xmax - xmin, ymax - ymin
    return (xmin + w/2, ymin + h/2, w, h)


def cxywh2xyxy(cxywh):
    """中心+宽高 -> 四方"""
    cx, cy, w, h = cxywh
    return (cx - w/2, cy - h/2, cx + w/2, cy + h/2)


def axywh2cxywh(axywh):
    """左上+宽高 -> 中心+宽高"""
    x, y, w, h = axywh
    return (x + w/2, y + h/2, w, h)


def cxywh2axywh(cxywh):
    """中心+宽高 -> 左上+宽高"""
    x, y, w, h = cxywh
    return (x - w/2, y - h/2, w, h)


def normalize_xywh(xywh, img_hw):
    """像素 -> 归一化"""
    x = float(xywh[0]) / img_hw[1]
    y = float(xywh[1]) / img_hw[0]
    w = float(xywh[2]) / img_hw[1]
    h = float(xywh[3]) / img_hw[0]
    return (x, y, w, h)


def reverse_normalized_xywh(nxywh, img_hw):
    """归一化 -> 像素"""
    # x, w = map(lambda x: int(x * iw), (x, w))
    # y, h = map(lambda y: int(y * ih), (y, h))
    x = int(nxywh[0] * img_hw[1])
    y = int(nxywh[1] * img_hw[0])
    w = int(nxywh[2] * img_hw[1])
    h = int(nxywh[3] * img_hw[0])
    return (x, y, w, h)


def try_as_nxywh(t4):
    """parse as float, only accepts range in [0.0, 1.0]"""
    f4 = list(map(float, t4))
    if not all(map(lambda x: 0.0 <= x and x <= 1.0, f4)):
        return float, None
    return float, f4


def try_as_cxywh(t4):
    """parse as int, also accepts i.0"""
    f4 = list(map(float, t4))
    i4 = list(map(int, t4))
    if not all(map(lambda i, f: abs(f - float(i)) < 0.001, zip(i4, f4))):
        return int, None
    return int, i4


# 这个算法有问题，w h theta 利用的参数不对
# def obb2cxywh(obb):
#     """四点 -近似-> 中心+宽高+角度。四点是（左上，右上，右下，左下）"""
#     x1, y1, x2, y2, x3, y3, x4, y4 = obb
#     cx = (x1 + x2 + x3 + x4) / 4
#     cy = (y1 + y2 + y3 + y4) / 4
#     w = math.sqrt((x1 - x2)**2 + (y1 - y2)**2)
#     h = math.sqrt((x2 - x3)**2 + (y2 - y3)**2)
#     theta = math.atan2(y2 - y1, x2 - x1)
#     return (cx, cy, w, h), theta


def cxywh2obb(cxywh, theta = 0.0):
    """中心+宽高+角度 -> 四点"""
    cx, cy, w, h = cxywh

    C = np.array([
        [cx, cy],
    ])

    R = np.array([
        [ math.cos(theta), math.sin(theta)],
        [-math.sin(theta), math.cos(theta)],
    ])

    pts = np.array([
        [-w/2, -h/2],
        [ w/2, -h/2],
        [ w/2,  h/2],
        [-w/2,  h/2],
    ])

    rot_pts = []
    for pt in pts:
        rot_pts.append((C + pt @ R))

    return np.array(rot_pts).flatten()


def obb2xyxy(obb):
    """四点 -> 四方，丢失旋转角。四方是（左上，右下）"""
    x1, y1, x2, y2, x3, y3, x4, y4 = obb
    xmin = min(x1, x2, x3, x4)
    xmax = max(x1, x2, x3, x4)
    ymin = min(y1, y2, y3, y4)
    ymax = max(y1, y2, y3, y4)
    return (xmin, ymin, xmax, ymax)


def xyxy2obb(xyxy, theta = 0.0):
    """四方+角度 -> 四方"""
    x1, y1, x2, y2 = xyxy

    # 四点
    p4 = np.array([
        [x1, y1],
        [x1, y2],
        [x2, y1],
        [x2, y2],
    ])

    # 旋转中心
    cx = np.mean(p4, axis=0)

    # 顺时针旋转
    transform = np.array([
        [math.cos(theta), -math.sin(theta)],
        [math.sin(theta),  math.cos(theta)],
    ])

    p4 = ((p4 - cx) @ transform) + cx
    return p4.flatten()


class ClassMixin:

    def __init__(self, classes_inout, classes_remap):
        self.classes = classes_inout
        self.classes_remap = classes_remap

    def _gset_class_id(self, name):
        if self.classes_remap:
            try:
                class_id = self.classes_remap[name]
                return class_id
            except ValueError:
                pass

        try:
            class_id = self.classes.index(name)
        except ValueError:
            self.classes.append(name)
            class_id = len(self.classes) - 1
        return class_id


class LabelTxtBase():

    @staticmethod
    def load_file(filename):
        boxes = []
        with open(filename, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if not line:
                    continue

                if line.startswith('#'):
                    continue

                parts = line.split()

                try:
                    if len(parts) >= 9:
                        # x1 y1 x2 y2 x3 y3 x4 y4 category:str [difficult:int]
                        obb = parts[:8]
                        category = parts[8]
                        boxes.append((category, obb, parts[9:]))
                        continue
                    elif len(parts) == 5 or len(parts) == 6:
                        # class_id:int, x0, y0, w0, h0, [angle]
                        class_id, x0, y0, w0, h0, *more = line.split()
                        class_id = int(class_id)
                        boxes.append((class_id, (x0, y0, w0, h0)))
                        continue
                    else:
                        raise ValueError('Unknown format: %s', line)
                except Exception as exc:
                    logger.warning('invalid box: %s', line, exc_info=exc)

        return boxes

    @staticmethod
    def dump_boxes(filename, boxes):
        with open(filename, 'w+') as f:
            for box in boxes:
                f.write(' '.join(box))
                f.write('\n')


class DotaObbTxt(LabelTxtBase, ClassMixin):

    def __init__(self, classes_inout, classes_remap, filename):
        LabelTxtBase.__init__(self)
        ClassMixin.__init__(self, classes_inout, classes_remap)
        self.filename = filename
        self.boxes = self.load_file(filename)

    def get_boxes(self, img_hw):
        boxes = []
        for category, obb, _ in self.boxes:
            xyxy = obb2xyxy(map(float, obb))
            cxywh = xyxy2cxywh(xyxy)
            class_id = self._gset_class_id(category)
            if class_id is None or class_id < 0:
                continue
            boxes.append((class_id, cxywh))
        return boxes


class YoloTxt(LabelTxtBase):

    def __init__(self, filename):
        super().__init__()
        self.filename = filename
        self.boxes = self.load_file(filename)

    def get_boxes(self, img_hw):
        boxes = []
        for class_id, (x0, y0, w0, h0), *more in self.boxes:
            dtype, nxywh = try_as_nxywh((x0, y0, w0, h0))
            if nxywh is None:
                logger.warning('invalid box: %s', str((x0, y0, w0, h0)))
                continue
            cxywh = reverse_normalized_xywh(nxywh, img_hw)
            boxes.append((class_id, cxywh))
        return boxes


class PlainTxt(LabelTxtBase):

    def __init__(self, filename):
        super().__init__()
        self.filename = filename
        self.boxes = self.load_file(filename)

    def get_boxes(self):
        boxes = []
        for class_id, (x0, y0, w0, h0) in self.boxes:
            dtype, cxywh = try_as_cxywh((x0, y0, w0, h0))
            if cxywh is None:
                logger.warning('invalid box: %s', str((x0, y0, w0, h0)))
                continue
            boxes.append((class_id, cxywh))
        return boxes


class LabelXml(ClassMixin):

    def __init__(self, classes_inout, classes_remap, filename):
        ClassMixin.__init__(self, classes_inout, classes_remap)
        self.filename = filename
        with open(filename, 'r') as f:
            root = lxml.etree.parse(f).getroot()
            if root.tag == 'annotation':
                self.root = root
            else:
                logger.error(lxml.etree.tostring(self.root, pretty_print=True))
                raise ValueError('Unknown root node: %s', root.tag)

    def get_image_filename(self):
        return self.root.xpath('source/filename')[0].text

    def get_chw(self):
        """channel, height, width"""
        node = self.root.xpath('size')
        if len(node) != 1:
            raise ValueError('len(node) != 1 : %s', len(node))
        node = node[0]

        img_w = int(node.xpath('width')[0].text)
        img_h = int(node.xpath('height')[0].text)
        img_c = int(node.xpath('depth')[0].text)
        return (img_c, img_h, img_w)

    def get_boxes(self, img_hw):
        boxes = []

        (_, img_h, img_w) = self.get_chw()
        if img_hw is not None:
            if img_hw[0] != img_h or img_hw[1] != img_w:
                raise ValueError('img_hw != (img_h, img_w) : %s != %s', img_hw, (img_h, img_w))
        else:
            img_hw = (img_h, img_w)

        objects = self.root.xpath('objects/object')
        for obj in objects:
            bndbox = obj.xpath('bndbox')
            if len(bndbox) > 0:
                # PASCAL VOC (Visual Object Challenge), labelImg
                assert len(bndbox) == 1
                bndbox = bndbox[0]

                name = obj.xpath('name')[0].text

                # 绝对值
                xmin = int(bndbox.xpath('xmin')[0].text)
                ymin = int(bndbox.xpath('ymin')[0].text)
                xmax = int(bndbox.xpath('xmax')[0].text)
                ymax = int(bndbox.xpath('ymax')[0].text)

                # 标注越界修正
                xmin, ymin = max(xmin, 0), max(ymin, 0)
                xmax, ymax = min(xmax, img_w), min(ymax, img_h)

                xyxy = (xmin, ymin, xmax, ymax)
                cxywh = xyxy2cxywh(xyxy)

                class_id = self._gset_class_id(name)
                if class_id is None or class_id < 0:
                    continue
                boxes.append((class_id, cxywh))
                continue

            robndbox = obj.xpath('robndbox')
            if len(robndbox) > 0:
                # roLabelImg TODO
                assert len(robndbox) == 1
                robndbox = robndbox[0]

                # 绝对值
                cx = float(robndbox.xpath('cx')[0].text)
                cy = float(robndbox.xpath('cy')[0].text)
                w = float(robndbox.xpath('w')[0].text)
                h = float(robndbox.xpath('h')[0].text)
                angle = float(robndbox.xpath('angle')[0].text)

                cxywh = (cx, cy, w, h)
                class_id = None
                boxes.append((class_id, cxywh, angle))
                continue

            if len(obj.xpath('points')) > 0:
                coordinate = obj.xpath('coordinate')[0].text
                if coordinate == 'pixel':
                    # FAIR1M
                    pass
                else:
                    logger.error('Unknown coordinate: %s', coordinate)
                    continue

                shape_type = obj.xpath('type')[0].text
                if shape_type == 'rectangle':
                    # FAIR1M
                    pass
                else:
                    logger.error('Unknown shape type: %s', shape_type)
                    continue

                names = obj.xpath('possibleresult/name')
                class_ids = filter(lambda x: not (x is None or x < 0), [self._gset_class_id(name.text) for name in names])
                if len(class_ids) != 1:
                    raise ValueError('len(class_ids) != 1')

                points = obj.xpath('points/point')
                if coordinate == 'pixel':
                    points = [list(map(float, p.text.split(','))) for p in points]

                if shape_type == 'rectangle':
                    if len(points) == 5:
                        if points[0] != points[4]:
                            raise ValueError('points[0] != points[4]')
                        points = points[:4]
                    elif len(points) == 4:
                        pass
                    else:
                        raise ValueError('len(points) != 4')

                    obb = flatten(points)
                    xyxy = obb2xyxy(obb)
                    cxywh = xyxy2cxywh(xyxy)
                    for class_id in class_ids:
                        boxes.append((class_id, cxywh))

                continue

            logger.error('Unknown xml format to parse objects')

        return boxes
